// Copyright 2022 jeff.li. and/or its affiliates.

#include <gtest/gtest.h>
#include <tbir/pipeline/pickle.h>
#include <tbir/runtime/generic/generic_constructor_funcs.h>
#include <tbir/runtime/json_util.h>

namespace tbir::runtime {

    TEST(Pipeline, pickle) {
        Dict d{{1,                 2},
               {Unicode(U"hello"), 3}};
        auto json_doc = pickle::ToJsonStruct(RTView(d));
        RTValue dd = pickle::FromJsonStruct(json_doc);
        ASSERT_EQ(dd, d);
        std::cout << JsonUtil::ToString(&json_doc, true) << std::endl;

        // Add Tupe
        auto tup = Tuple::dynamic(1, 2.5f, "a", List({1, 2}));
        auto tup_json = pickle::ToJsonStruct(RTView(tup));
        auto new_tup = pickle::FromJsonStruct(tup_json).AsObjectRefNoCheck<Tuple>();
        ASSERT_EQ(tup.size(), new_tup.size());
        ASSERT_EQ(tup[3], new_tup[3]);
    }

    TEST(Pickle, ndarray) {
        auto arr = ::tbir::runtime::Kernel_NDArray::make({1.0, 2.0, 3.0, 4.0}, {2, 2}, U"float32");
        auto json_doc = pickle::ToJsonStruct(RTView(arr));

        auto new_arr = pickle::FromJsonStruct(json_doc).AsObjectRefNoCheck<NDArray>();
        ASSERT_EQ(arr.ShapeList(), new_arr.ShapeList());
        ASSERT_EQ(arr.DTypeUnicode(), new_arr.DTypeUnicode());
    }

    TEST(Pickle, serialzation) {
        auto arr = ::tbir::runtime::Kernel_NDArray::make({1.0, 2.0, 3.0, 4.0}, {2, 2}, U"float32");
        auto json_doc = pickle::ToJsonStruct(RTView(arr));

        auto str = pickle::Serialize(RTView(arr));
        auto new_arr = pickle::DeSerialize(str);
        auto shape = new_arr.AsObjectRef<NDArray>().Shape();
        ASSERT_EQ(shape[0], 2);
        ASSERT_EQ(shape[1], 2);
    }

}  // namespace tbir::runtime
